#coding:utf8
"""
至少5天数据才可以提取对应特征进行预测
获取当天日期，利用code获取对应文件，判断日期，获取前四天的数据，提取特征
具体方式，首先全部获取，然后获取前几天数据
修改时间：20220903
对应数据：get_train_data_02
"""
import os, re, json,time
import datetime
import akshare as ak
import pandas as pd
import joblib
import numpy as np


def isEffective(days, all_trade_dates):
    index_today = all_trade_dates.index(days[0])
    flag = True
    for i in range(1,len(days)):
        if days[i] != all_trade_dates[index_today - i]:
            flag = False
            break

    return flag


def zjlx_change(z1, z2):
    f105 = 0
    if np.sign(z1) > 0:
        if np.sign(z2) <= 0:
            f105 = 1
        elif z1 > z2:
            f105 = 2
        else:
            f105 = 3
    else:
        if np.sign(z2) >= 0:
            f105 = 4
        elif abs(z1) > abs(z2):
            f105 = 5
        else:
            f105 = 6

    return f105


def get_latest_5_days(date=None):
    #获取所有交易日
    tool_trade_date = ak.tool_trade_date_hist_sina()
    tmp1 = tool_trade_date['trade_date'].tolist()
    all_trade_dates = [m.strftime("%Y-%m-%d") for m in tmp1]

    #获取当天日期
    if not date:
        now = datetime.datetime.now()
        td = now.strftime("%Y-%m-%d")
    else:
        td = date

    latest5days = []
    if td in all_trade_dates:
        print(f"当天是交易日")
        _id = all_trade_dates.index(td)
        for i in range(_id, _id -  5, -1):
            latest5days.append(all_trade_dates[i])
    else:
        print(f"当天不是交易日")
        for i in range(1,30):
            d1 = datetime.datetime.strptime(td, "%Y-%m-%d")
            d2 = d1 + datetime.timedelta(days=-i)
            d3 = datetime.datetime.strftime(d2, "%Y-%m-%d")
            if d3 in all_trade_dates:
                _id = all_trade_dates.index(d3)
                break
        for i in range(_id, _id - 5, -1):
            d1 = all_trade_dates[i]
            latest5days.append(d1)


    print(f"最近5个交易日(含当天)：{latest5days}")
    return latest5days


def get_feature_data_01(df, infer=False):
    Y, X = [], []
    try:
        i = 0
        #资金流入
        tmp_1 = df['主力净流入'].astype(float)
        tmp_12 = df['流通市值'].astype(float)
        f11 = int(100 * tmp_1[i] / tmp_12[i])
        f12 = 1 if tmp_1[i] > 0 else 0
        tmp_13 = tmp_1[i:i + 5].sum()  # 5日累计主力净流入
        f13 = int(100 * tmp_13 / tmp_12[i])
        f14 = 1 if tmp_13 > 0 else 0
        tmp_13 = tmp_1[i:i + 3].sum()  # 3日累计主力净流入
        f15 = int(100 * tmp_13 / tmp_12[i])
        f16 = 1 if tmp_13 > 0 else 0
        #资金流入增幅
        f17 = 1 if tmp_1[i] - tmp_1[i+1] else 0
        f18 = 1 if tmp_1[i+1] - tmp_1[i+2] else 0
        f19 = 1 if tmp_1[i+1] > 0 else 0    #前一天

        #成交量
        tmp_2 = df['成交量'].astype(float)
        tmp_21 = df['成交额'].astype(float)
        tmp_22 = tmp_2[i] + tmp_2[i + 2] - 2 * tmp_2[i + 1]
        f21 = int(100 * tmp_21[i] / tmp_12[i])
        f22 = int(tmp_21[i:i+5].sum() / 5 / tmp_12[i])
        f23 = int(tmp_21[i:i+3].sum() / 3 / tmp_12[i])
        f24 = 1 if tmp_2[i] - tmp_2[i+1] > 0 else 0
        f25 = 1 if tmp_2[i+1] - tmp_2[i+2] > 0 else 0
        f26 = 1 if tmp_22 > 0 else 0

        #涨跌幅
        tmp_3 = df['涨跌幅'].astype(float)
        tmp_32 = tmp_3[i:i + 5].sum()  # 5日累计涨跌幅
        tmp_33 = tmp_3[i:i + 3].sum()  # 3日累计涨跌幅
        cnt1 = 0
        for i in range(5):
            if tmp_3[i] >= 0:
                cnt1 += 1
            else:
                break
        f31 = 1 if tmp_3[i] > 0 else 0
        f32 = 1 if tmp_3[i+1] > 0 else 0        #前一天
        f33 = 1 if tmp_32 > 0 else 0
        f34 = 1 if tmp_33 > 0 else 0
        f35 = cnt1
        #f35 =

        #
        tmp_4 = df['换手率'].astype(float)
        tmp_42 = tmp_4[i] - tmp_4[i + 1]  # 收盘换手率增幅
        tmp_43 = tmp_4[i] + tmp_4[i + 2] - 2 * tmp_4[i + 1]  # 收盘换手率增速
        f41 = tmp_4[i]
        f42 = 1 if tmp_42 > 0 else 0
        f43 = 1 if tmp_4[i+1] - tmp_4[i+2] > 0 else 0    #前一天换手是否增加
        f43 = 1 if tmp_43 > 0 else 0        #增加幅度是否增加
        f44 = tmp_4[i:i + 5].sum() / 5      #5日平均换手率
        f45 = (f41 - f44) / (f41 + 1e-4)

        #以上所有因子暂时未与股价挂钩，故增加股价因子
        tmp_5 = df['收盘价'].astype(float)
        tmp_52 = tmp_5[i:i + 5].sum()
        tmp_53 = tmp_5[i:i + 3].sum()
        f51 = tmp_5[i]
        f52 = (f51 - tmp_52 / 5) / (f51 + 1e-4)
        f53 = (f51 - tmp_53 / 3) / (f51 + 1e-4)
        # f54 = df['开盘价'] - df['']

        # item = ['主力净流入', '5日累计主力净流入', '3日累计主力净流入', '涨幅', '5日累计涨幅', '3日累计涨幅', '成交量增幅', '成交量增速', '换手率', '收盘换手率增幅',
        #         '收盘换手率增速']
        # x = [df['主力净流入'].iloc[i], tmp_12, tmp_13, df['涨跌幅_x'].iloc[i], tmp_22, tmp_23, tmp_32, tmp_33,
        #      df['换手率'].iloc[i], tmp_42, tmp_43]
        x = [f11, f12, f13, f14, f15, f16, f17, f18, f19,
             f21, f22, f23, f24, f25, f26,
             f31, f32, f33, f34, f35,
             f41, f42, f43, f44, f45,
             f51,f52,f53]

        X.append(x)
        if not infer:
            y = 1 if df['涨跌幅'].iloc[i - 1] > 3 else 0
            Y.append(y)
    except Exception as e:
        print(e.__str__())

    return Y, X


def get_feature_data_02(df, all_trade_dates,savePath=None):
    X, Y = [],[]

    i = 0
    today = df['日期'][i:i+5].values
    flag = isEffective(today, all_trade_dates)
    if not flag:
        return None

    #资金流入
    tmp_11 = df['主力净流入'].astype(float)
    tmp_12 = df['流通市值'].astype(float)
    tmp_13 = df['收盘价'].astype(float)
    tmp_14 = df['主力净流入占比'].astype(float)
    f101 = 1 if tmp_11[i] > 0 else 0
    f102 = 100 * tmp_11[i] / tmp_12[i]     #主力流入/流通市值
    f103 = tmp_14[i]                 #主力流入占比
    f104 = np.sign(tmp_11[i]) * np.sqrt(abs(tmp_11[i]/10000) * tmp_13[i])  #主力流入价格指数
    #累计资金流入
    tmp_15 = tmp_11[i:i + 5].sum()  # 5日累计主力净流入
    tmp_16 = tmp_11[i:i + 3].sum()  # 3日累计主力净流入
    f105 = 1 if tmp_15 > 0 else 0
    f106 = 100 * tmp_15 / tmp_12[i]
    f107 = np.sign(tmp_15) * np.sqrt(abs(tmp_15/10000) * tmp_13[i])
    f108 = 1 if tmp_16 > 0 else 0
    f109 = 100 * tmp_16 / tmp_12[i]
    f110 = np.sign(tmp_16) * np.sqrt(abs(tmp_16/10000) * tmp_13[i])
    #资金流入变动
    f111 = 100 * np.sign(tmp_11[i]) * (tmp_11[i] - tmp_11[i+1]) / tmp_12[i]
    f112 = 100 * np.sign(tmp_11[i+1]) * (tmp_11[i+1] - tmp_11[i + 2]) / tmp_12[i+1]
    f113 = 100 * np.sign(tmp_11[i+2]) * (tmp_11[i+2] - tmp_11[i + 3]) / tmp_12[i+2]
    #f109 =

    #
    tmp_51 = df['收盘价'].astype(float)
    tmp_52 = df['今开'].astype(float)
    tmp_53 = df['昨日收盘'].astype(float)
    #tmp_54 = df['收盘价_x'].astype(float)
    tmp_55 = df['涨跌幅'].astype(float)
    tmp_56 = df['最高'].astype(float)
    tmp_57 = df['最低'].astype(float)

    f206 = 1 if tmp_55[i] > 0 else 0
    f207 = 1 if tmp_52[i] > tmp_53[i] else 0
    f208 = 1 if tmp_52[i] >= tmp_57[i] else 0
    f209 = 1 if tmp_51[i] >= tmp_56[i] else 0
    f210 = 1 if tmp_51[i] > tmp_51[i+1] else 0
    f211 = 1 if tmp_51[i+1] > tmp_51[i+2] else 0
    f212 = 1 if tmp_51[i] > np.mean(tmp_51[i:i+3]) else 0
    f213 = 1 if tmp_51[i] > np.mean(tmp_51[i:i+5]) else 0

    tmp_21 = df['涨跌幅'].astype(float)
    tmp_22 = tmp_21[i:i + 5].sum()  # 5日累计涨跌幅
    tmp_23 = tmp_21[i:i + 3].sum()  # 3日累计涨跌幅
    f201 = tmp_21[i]
    f202 = tmp_22
    f203 = tmp_23
    f204 = tmp_21[i+1]
    f205 = tmp_21[i] - tmp_21[i+1]

    #
    tmp_31 = df['成交量'].astype(float)
    tmp_32 = df['成交金额'].astype(float)
    f301 = tmp_32[i] / tmp_12[i]  # 成交金额/流通市值
    f302 = np.sqrt(tmp_32[i] * tmp_13[i]) #成交价格指数
    f303 = (tmp_31[i] - tmp_31[i+1]) / tmp_31[i+1]  #成交量增加率
    f304 = (tmp_31[i+1] - tmp_31[i + 2]) / tmp_31[i + 2]  # 成交量增加率
    f305 = (tmp_31[i+2] - tmp_31[i + 3]) / tmp_31[i + 3]  # 成交量增加率
    f306 = tmp_31[i] + tmp_31[i+2] - 2 * tmp_31[i+1]  # 成交量增速

    f307 = 1 if tmp_31[i] > tmp_31[i+1] else 0
    f308 = 1 if tmp_31[i+1] > tmp_31[i+2] else 0
    f309 = 1 if tmp_31[i] > np.mean(tmp_31[i:i+3]) else 0
    f310 = 1 if tmp_31[i] > np.mean(tmp_31[i:i+5]) else 0

    #
    tmp_41 = df['换手率'].astype(float)
    f401 = tmp_41[i]
    f402 = (tmp_41[i] - tmp_41[i+1])/(tmp_41[i+1] + 1e-3) # 收盘换手率增幅
    f403 = (tmp_41[i+1] - tmp_41[i+2])/(tmp_41[i+2] + 1e-3) # 收盘换手率增幅
    f404 = (tmp_41[i+2] - tmp_41[i+3])/(tmp_41[i+3] + 1e-3) # 收盘换手率增幅
    f405 = np.sqrt(tmp_41[i] * tmp_13[i]) #换手率价格指数
    f406 = tmp_41[i] + tmp_41[i+2] - 2 * tmp_41[i+1]  # 收盘换手率增速

    f407 = 1 if tmp_41[i] > 0 else 0
    f408 = 1 if tmp_41[i] > tmp_41[i+1] else 0
    f409 = 1 if tmp_41[i] > np.mean(tmp_41[i:i+3]) else 0
    f410 = 1 if tmp_41[i] > np.mean(tmp_41[i:i + 5]) else 0

    x = [f101, f102, f103, f104,f105,f106,f107,f108,f109,f110,f111,f112,f113,
         f201,f202,f203,f204,f205,f206,f207,f208,f209,f210,f211,f212,f213,
         f301,f302,f303,f304,f305,f306,f307,f308,f309,f310,
         f401,f402,f403,f404,f405,f406,f407,f408,f409,f410]
    #y = 1 if df['涨跌幅'].iloc[i - 1] > 3 else 0
    X.append(x)
    #Y.append(y)

    return Y,X


def get_feature_data_03(df, all_trade_dates,savePath=None):
    X, Y = [],[]

    i = 0
    today = df['日期'][i:i+5].values
    flag = isEffective(today, all_trade_dates)
    if not flag:
        return None

    #资金流入
    tmp_11 = df['主力净流入'].astype(float)
    tmp_12 = df['流通市值'].astype(float)
    tmp_13 = df['收盘价'].astype(float)
    tmp_14 = df['主力净流入占比'].astype(float)
    tmp_15 = tmp_11[i:i + 5].sum()  # 5日累计主力净流入
    tmp_16 = tmp_11[i:i + 3].sum()  # 3日累计主力净流入

    # 主力净流入
    f101 = 1 if tmp_11[i] > 0 else 0    #当天主力净流入
    f102 = 1 if tmp_11[i+1] > 0 else 0
    f103 = 1 if tmp_16 > 0 else 0
    f104 = 1 if tmp_15 > 0 else 0

    # 主力净流入变化
    f105 = zjlx_change(tmp_11[i], tmp_11[i+1])
    f106 = zjlx_change(tmp_11[i], tmp_16)               # 当天与近三天
    f107 = zjlx_change(tmp_11[i], tmp_15)               # 当天与近五天

    #成交量
    tmp_31 = df['成交量'].astype(float)
    tmp_32 = df['成交金额'].astype(float)
    f201 = 1 if tmp_31[i] > tmp_31[i+1] else 0
    f202 = 1 if tmp_31[i+1] > tmp_31[i+2] else 0
    f203 = 1 if tmp_31[i] > np.mean(tmp_31[i:i+3]) else 0
    f204 = 1 if tmp_31[i] > np.mean(tmp_31[i:i+5]) else 0

    #换手率
    tmp_41 = df['换手率'].astype(float)
    f301 = 1 if tmp_41[i] > 0 else 0
    f302 = 1 if tmp_41[i] > tmp_41[i+1] else 0
    f303 = 1 if tmp_41[i] > np.mean(tmp_41[i:i+3]) else 0
    f304 = 1 if tmp_41[i] > np.mean(tmp_41[i:i + 5]) else 0

    #涨幅
    tmp_51 = df['收盘价'].astype(float)
    tmp_52 = df['今开'].astype(float)
    tmp_53 = df['昨日收盘'].astype(float)
    #tmp_54 = df['收盘价_x'].astype(float)
    tmp_55 = df['涨跌幅'].astype(float)
    tmp_56 = df['最高'].astype(float)
    tmp_57 = df['最低'].astype(float)

    f401 = 1 if tmp_55[i] > 0 else 0
    f402 = 1 if tmp_52[i] > tmp_53[i] else 0
    f403 = 1 if tmp_52[i] >= tmp_57[i] else 0
    f404 = 1 if tmp_51[i] >= tmp_56[i] else 0
    f405 = 1 if tmp_51[i] > tmp_51[i+1] else 0
    f406 = 1 if tmp_51[i+1] > tmp_51[i+2] else 0
    f407 = 1 if tmp_51[i] > np.mean(tmp_51[i:i+3]) else 0
    f408 = 1 if tmp_51[i] > np.mean(tmp_51[i:i+5]) else 0

    item = ['主力净流入', '5日累计主力净流入', '3日累计主力净流入', '涨幅', '5日累计涨幅', '3日累计涨幅', '成交量增幅', '成交量增速', '换手率', '收盘换手率增幅',
            '收盘换手率增速']

    x = [f101, f102, f103, f104,f105,f106,f107,
         f201,f202,f203,f204,
         f301,f302,f303,f304,
         f401,f402,f403,f404,f405,f406,f407,f408]
    #y = 1 if df['涨跌幅'].iloc[i - 1] > 3 else 0
    X.append(x)
    #Y.append(y)

    return Y,X


def get_data(date=None):
    """
    获取realtime_data 中的数据
    :return:
    """

    tool_trade_date = ak.tool_trade_date_hist_sina()
    tmp1 = tool_trade_date['trade_date'].tolist()
    all_trade_dates = [m.strftime("%Y-%m-%d") for m in tmp1]

    #获取前5天交易日期,日期由近及远排列
    latest5date = get_latest_5_days(date=date)
    #latest5date = ['2022-09-08', '2022-09-07', '2022-09-06', '2022-09-05', '2022-09-02']
    #latest5date = ['2022-09-15', '2022-09-14', '2022-09-13', '2022-09-09', '2022-09-08']
    #首先从实时数据中查找
    files = []
    for m in latest5date:
        fileName = m + ".csv"
        filePath = os.path.join(r"../realtime_data", fileName)
        if os.path.exists(filePath):
            files.append(filePath)


    if len(files) == 5:
        print(r"最近五天数据都存在")
        df_total = pd.DataFrame(data=None, columns=None)
        dfs = []
        for file in files:
            df_tmp = pd.read_csv(file, dtype={'股票代码': str})
            dfs.append(df_tmp)

        df_total = pd.concat(dfs, ignore_index=True)
        del dfs
        # 获取所有股票的5天数据
        df_result = pd.DataFrame(data=None, columns=None)
        if files:
            stock_tmp = pd.read_csv(files[0], dtype={'股票代码': str})['股票代码'].tolist()
            stock_list = [m for m in stock_tmp if isinstance(m, str) and re.findall(r"\d{6}", m)]
        else:
            stock_list = list(df_total['股票代码'].values)
            pass
        cnt = 0
        feature_data = {}
        for code in stock_list:
            try:
                if re.findall(r"^4|^3|^8|^688", code):
                    continue
                print(f"extract data of: {code}")
                df_tmp = pd.DataFrame(data=None, columns=None)
                dfs = []
                for day in latest5date:
                    t2 = df_total[(df_total['股票代码'] == code) & (df_total['日期'] == day)]
                    t21 = df_total.loc[df_total['股票代码'] == code]
                    t22 = df_total.loc[df_total['日期'] == day]
                    if t2.empty or t2['股票代码'].isnull().any() or t2['股票名称'].isnull().any():
                        break
                    dfs.append(t2)
                df_tmp = pd.concat(dfs, ignore_index=True)
                df_tmp = df_tmp.reset_index(drop=True)

                # print(df_tmp)
                if df_tmp.empty or len(df_tmp) < 5:
                    continue
                else:
                    features = get_feature_data_02(df_tmp, infer=True)
                    if not features:
                        continue
                    elif features[1]:
                        feature_data[code] = features[1]
                        cnt += 1
                        print(f" 有效数据: {cnt}")
            except Exception as e:
                print(f"error: {code}")

            #break

    else:
        print(f"最近5天数据不够，需从历史数据补充")
        #获取所有数据
        df_total = pd.DataFrame(data=None, columns=None)
        dfs = []
        for i in range(5):
            if i < len(files):
                dt = latest5date[i]
                df = pd.read_csv(files[i], dtype={'股票代码': str})
                dfs.append(df)
                #df_total = df_total.append(df, ignore_index=True)
            else:
                history_dir = r"../data/history/20221022"
                for f in os.listdir(history_dir):
                    if re.findall(r"^4|^3|^8|^688", f):
                        continue
                    # if not re.findall(r"000090", f):
                    #     continue
                    p1 = os.path.join(history_dir, f)
                    t1 = pd.read_csv(p1)
                    t1.rename(columns={"股票代码_x":"股票代码", "涨跌幅_x":"涨跌幅", "成交额":"成交金额", "收盘价_x":"收盘价", "开盘价":"今开",
                                       "最高价":"最高", "最低价":"最低","前收盘":"昨日收盘"}, inplace=True)
                    #t1['股票代码'] = t1['股票代码'].astype(str)
                    if t1.empty:
                        print(f"{f} 文件为空")
                        continue

                    for j in range(i, 5):
                        dt = latest5date[j]
                        t2 = t1[t1['日期'] == dt]
                        if t2.empty:
                            break
                        #t2 = t2.reset_index(drop=True)
                        #t3 = t2['股票代码']
                        #t4 = t2['股票代码'].iloc[0]
                        if len(t2['股票代码'].iloc[0]) == 7:
                            t2['股票代码'].iloc[0] = t2['股票代码'].iloc[0][1:]
                        #df_total = df_total.append(t2, ignore_index=True)
                        t5 = isinstance(t2, pd.DataFrame)
                        dfs.append(t2)
                    break
                break


        df_total = pd.concat(dfs, ignore_index=True)
        #获取所有股票的5天数据
        df_result = pd.DataFrame(data=None, columns=None)
        stock_list = set(list(df_total['股票代码'].values))
        #t1 = pd.read_csv(files[0], dtype={'股票代码': str})
        #stock_tmp = t1['股票代码'].tolist()
        #stock_list = [m for m in stock_tmp if isinstance(m, str) and re.findall(r"\d{6}", m)]
        cnt = 0
        feature_data = {}
        for code in stock_list:
            try:
                if re.findall(r"^3|^8|^688", code):
                    continue
                # if not re.findall(r"003037", code):
                #     continue
                print(f"extract data of: {code}")
                df_tmp = pd.DataFrame(data=None, columns=None)
                for day in latest5date:
                    t2 = df_total[(df_total['股票代码'] == code) & (df_total['日期'] == day)]
                    #t21 = df_total.loc[df_total['股票代码'] == code]
                    #t22 = df_total.loc[df_total['日期'] == day]
                    if t2.empty or t2['股票代码'].isnull().any() or t2['股票名称'].isnull().any():
                        break
                    df_tmp = df_tmp.append(t2, ignore_index=True)
                    df_tmp = df_tmp.reset_index(drop=True)

                #print(df_tmp)
                if df_tmp.empty or len(df_tmp) < 5:
                    continue
                else:
                    features = get_feature_data_02(df_tmp, all_trade_dates)
                    if features[1]:
                        feature_data[code] = features[1]
                        cnt += 1
                        print(f" 有效数据: {cnt}")
            except Exception as e:
                print(f"error: {code}")
                print(e.__str__())
            #break
        print(len(feature_data))


    return feature_data


def inference(modelPath, date=None):
    time_0 = time.time()
    #获取数据
    feature_data = get_data(date=date)
    print(f"load data time cost: {time.time() - time_0}")

    #模型预测
    time_1 = time.time()
    #filePath = "../resources/20220825_05_rf.m"
    model = joblib.load(modelPath)
    print("="*100)
    result = []
    codes, X = [], []
    for code in feature_data:
        codes.append(code)
        X.append(feature_data[code][0])
    p2 = model.predict_proba(X)
    time_2 = time.time()
    print(f"model time cost: {time_2 - time_1}")

    for i in range(len(codes)):
        if p2[i][1] > 0.5:
            result.append((codes[i], p2[i][1]))

    t1 = sorted(result, key=lambda x:x[1], reverse=True)
    with open("result_0905.txt", "w", encoding="utf8") as f:
        for line in t1:
            print(line)
            tmp_str = str(line[0]) + "\t" + str(line[1])
            f.write(tmp_str + "\n")


if __name__=="__main__":
    modelPath = "../resources/20221022_02_rf.m"
    date = '2022-10-20'
    date = None
    inference(modelPath=modelPath, date=date)
    #get_data()
